﻿#include "PPMatting.h"

#ifdef WIN32
const char sep = '\\';
#else
const char sep = '/';
#endif


cv::Mat MattingInfer(fastdeploy::vision::matting::PPMatting& model, const cv::Mat image, const std::string& background_file = "") {
    cv::Mat vis_im;
    auto im = image;
    fastdeploy::vision::MattingResult res;
    if (!model.Predict(&im, &res)) {
        std::cerr << "Failed to predict." << std::endl;
        return vis_im;
    }

    if (!background_file.empty()) {
        auto bg = cv::imread(background_file);
        vis_im = fastdeploy::vision::SwapBackground(im, bg, res);
    }
    else {
        vis_im = fastdeploy::vision::VisMatting(im, res);
    }
    return vis_im;
}

fastdeploy::vision::matting::PPMatting CpuInit(const std::string& model_dir) {
    auto model_file = model_dir + sep + "model.pdmodel";
    auto params_file = model_dir + sep + "model.pdiparams";
    auto config_file = model_dir + sep + "deploy.yaml";
    auto option = fastdeploy::RuntimeOption();
    option.UseCpu();
    auto model = fastdeploy::vision::matting::PPMatting(model_file, params_file,
        config_file, option);
    if (!model.Initialized()) {
        std::cerr << "Failed to initialize." << std::endl;
        //return;
    }
    return model;
}

fastdeploy::vision::matting::PPMatting GpuInit(const std::string& model_dir) {
    auto model_file = model_dir + sep + "model.pdmodel";
    auto params_file = model_dir + sep + "model.pdiparams";
    auto config_file = model_dir + sep + "deploy.yaml";
    auto option = fastdeploy::RuntimeOption();
    option.UseGpu();
    option.UsePaddleInferBackend();
    auto model = fastdeploy::vision::matting::PPMatting(model_file, params_file,
        config_file, option);
    if (!model.Initialized()) {
        std::cerr << "Failed to initialize." << std::endl;
    }
    return model;
}
